package cn.edu.flink.scala.tutorial.state.keyedstate

import cn.edu.flink.scala.tutorial.source.{ClickSource, Event}
import jdk.jfr.Timestamp
import org.apache.flink.api.common.functions.{AggregateFunction, FlatMapFunction, RichFlatMapFunction}
import org.apache.flink.api.common.state.{AggregatingState, AggregatingStateDescriptor, ValueState, ValueStateDescriptor}
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

object AggregatingStateTest {
  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI()

    val clickStream = env.addSource(new ClickSource).assignAscendingTimestamps(_.timestamp)

    clickStream.keyBy(_.user)
      .flatMap(new RichFlatMapFunction[Event, String] {

        // 统计用户点击5c次的平均时间戳
        lazy val aggregatingState: AggregatingState[Event, Long] = getRuntimeContext.getAggregatingState(new AggregatingStateDescriptor[Event, (Long, Long), Long]("count state",
          new AggregateFunction[Event, (Long, Long), Long] {
            override def createAccumulator(): (Long, Long) = (0L, 0L)

            override def add(value: Event, accumulator: (Long, Long)): (Long, Long) = (accumulator._1 + value.timestamp, accumulator._2 + 1)

            override def getResult(accumulator: (Long, Long)): Long = accumulator._1 / accumulator._2

            override def merge(a: (Long, Long), b: (Long, Long)): (Long, Long) = ???
          },
          classOf[(Long, Long)]))

        lazy val numState: ValueState[Int] = getRuntimeContext.getState(new ValueStateDescriptor[Int]("num state", classOf[Int]))

        override def flatMap(value: Event, out: Collector[String]): Unit = {
          aggregatingState.add(value)

          numState.update(numState.value()+1)

          if (numState.value() == 5) {
            out.collect(s"${value.user} 完成5次点击，5次点击平均时间戳为：${aggregatingState.get()}")
//            aggregatingState.clear()
            numState.clear()
          }
        }
      }).print()

    env.execute("AggregatingStateTest")
  }
}
